clear all;
clc;

load mnist_train;

Y = Yl;
X = Xl;
clear Xl;
clear Yl;
k = 40;
[n,d] = size(X);
cls = unique(Y);
c = size(cls,1);

nn = zeros(n,k);
dt = zeros(n,k);
for i = 1:c
    i
    idx = find(Y==cls(i));
    x = X(idx,:);
    [nx,dx] = knnsearch(x,x,'K',k+1,'Distance','cosine');
    nn(idx,:) = idx(nx(:,2:k+1));
    dt(idx,:) = dx(:,2:k+1);
end

load mnist_test;
clear Yu;
[tnn,tdt] = knnsearch(X,Xu,'K',k,'Distance','cosine');

nn = [nn;tnn];
dt = [dt;tdt];

save('mnist_knn_cos_sup','nn','dt');
